import requests
import json
import os
import re
import pandas as pd
import time
from datetime import datetime

def query_dify(question):
    """向Dify API发送请求并获取回答"""
    url = "http://dify.sinohealth.cn/v1/chat-messages"
    
    payload = json.dumps({
       "inputs": {},
       "query": question,
       "response_mode": "streaming",
       "conversation_id": "",
       "user": "abc-123",
       "files": [
          {
             "type": "image",
             "transfer_method": "remote_url",
             "url": "https://cloud.dify.ai/logo/logo-site.png"
          }
       ]
    })
    
    headers = {
       'Authorization': 'Bearer app-D4ZielrNkZRr0AS7VW1cnApr',
       'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
       'Content-Type': 'application/json; charset=utf-8',
    }
    
    response = requests.request("POST", url, headers=headers, data=payload)
    
    # 确保使用正确的编码
    response.encoding = 'utf-8'
    
    # 获取响应内容
    response_text = response.text
    
    # 处理响应内容
    formatted_response = ""
    answer_content = ""
    lines = response_text.split('\n')
    
    for line in lines:
        # 对于每一行，先尝试解析为JSON
        if line.startswith('data: '):
            try:
                # 提取JSON部分
                json_str = line[6:]  # 去掉 'data: ' 前缀
                json_obj = json.loads(json_str)
                
                # 检查是否包含最终答案（在workflow_run_id.outputs.answer中）
                if 'data' in json_obj and 'outputs' in json_obj.get('data', {}) and 'answer' in json_obj['data'].get('outputs', {}):
                    answer_content = json_obj['data']['outputs']['answer']
                
                # 将JSON转换为字符串，确保不使用ASCII编码
                json_str = json.dumps(json_obj, ensure_ascii=False, indent=2)
                # 添加到格式化响应中
                formatted_response += "data: " + json_str + "\n"
            except json.JSONDecodeError:
                # 如果不是有效的JSON，直接添加原始行
                formatted_response += decode_unicode_in_text(line) + "\n"
        else:
            # 非JSON行，直接添加，但解码Unicode转义序列
            formatted_response += decode_unicode_in_text(line) + "\n"
    
    # 保存响应内容到文件
    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"dify_response_{current_time}.txt"
    
    with open(filename, 'w', encoding='utf-8') as file:
        file.write(f"查询: {question}\n\n")
        file.write(f"响应时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
        file.write("响应内容:\n")
        file.write(formatted_response)
    
    # 尝试提取表格数据
    table_data = ""
    try:
        # 尝试使用正则表达式提取"最终数据"部分
        table_match = re.findall(r'- 最终数据：(.*?)- 数据解读：', answer_content, re.DOTALL)
        if table_match:
            table_data = table_match[0].strip()
    except Exception as e:
        print(f"提取表格数据时出错: {e}")
        table_data = "提取失败"
    
    return answer_content, table_data, formatted_response

# 定义一个函数来处理Unicode转义序列
def decode_unicode_in_text(text):
    # 使用正则表达式查找所有形如 \uXXXX 的Unicode转义序列并解码
    return re.sub(r'\\u([0-9a-fA-F]{4})', lambda m: chr(int(m.group(1), 16)), text)

# 主函数
def main():
    # 读取问题测试集
    try:
        test_df = pd.read_excel("问题测试集.xlsx")
        print(f"成功读取测试集，共有 {len(test_df)} 个问题")
    except Exception as e:
        print(f"读取问题测试集时出错: {e}")
        return
    
    # 确保DataFrame中有必要的列
    required_columns = ['问题', '预期结果']
    for col in required_columns:
        if col not in test_df.columns:
            print(f"错误: 测试集中缺少必要的列 '{col}'")
            return
    
    # 添加结果列
    if '实际结果' not in test_df.columns:
        test_df['实际结果'] = ""
    if '完整回答' not in test_df.columns:
        test_df['完整回答'] = ""
    
    # 遍历问题并查询
    for index, row in test_df.iterrows():
        question = row['问题']
        print(f"\n处理问题 {index+1}/{len(test_df)}: {question}")
        
        try:
            # 查询Dify API
            answer_content, table_data, formatted_response = query_dify(question)
            
            # 更新DataFrame
            test_df.at[index, '实际结果'] = table_data
            test_df.at[index, '完整回答'] = answer_content
            
            print(f"已获取回答，表格数据长度: {len(table_data)}")
            
            # 每次查询后暂停几秒，避免请求过于频繁
            time.sleep(3)
        except Exception as e:
            print(f"处理问题时出错: {e}")
            test_df.at[index, '实际结果'] = f"错误: {str(e)}"
    
    # 保存结果到新的Excel文件
    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_filename = f"问题测试结果_{current_time}.xlsx"
    test_df.to_excel(result_filename, index=False, engine='openpyxl')
    
    print(f"\n所有问题处理完毕，结果已保存到: {os.path.abspath(result_filename)}")

# 执行主函数
if __name__ == "__main__":
    main()
